/*
  Copyright 2005 - Daniel Lowd

  This file is part of the nbe code available at http://www.cs.washington.edu/ai/nbe/. There are some modifications to 
  integrate the original code of Daniel Lowd into the eapmlib library. The NBE method is explained in:

  Daniel Lowd and Pedro Domingos. 'Naive Bayes Models for Probability Estimation'. 
  Proceedings of the Twenty-Second International Conference on Machine Learning (ICML), 2005. 
  Bonn, Germany: ACM Press.
*/
/*!
\file AbsVarAbstraction.h
\brief 
\author Daniel Lowd

*/
#ifndef ABSVARABSTRACTIONSET_H
#define ABSVARABSTRACTIONSET_H

#include "EvolutiveLib.h"
#include "AbstractionSet.h"

namespace NBE {

#define M_PRIOR 0.1
#define INITIAL_WEIGHT 1.0

// HACK
#define TEST_HACKS

class AbsVarAbstractionSet : public AbstractionSet
{
    double**  globalVarValProbs;
    bool fuzzyAbs;
    int  numObs;
    double vardiv;

public:
    AbsVarAbstractionSet(bool fuzzy = false, double div = 3.0)
        :AbstractionSet(), globalVarValProbs(0), fuzzyAbs(fuzzy), vardiv(div)
            {}

    AbsVarAbstractionSet(const VarSchema& s, bool fuzzy = false, 
            double div = 3.0) 
        :AbstractionSet(s), globalVarValProbs(0), fuzzyAbs(fuzzy), vardiv(div)
            {}

    AbsVarAbstractionSet(const AbsVarAbstractionSet& other);
    AbsVarAbstractionSet& operator=(const AbsVarAbstractionSet& other);

    virtual ~AbsVarAbstractionSet();

    // TODO: add support for continuous vars in this
    virtual VarSet getSample(int a) const;

    virtual Prob getProb(const VarSet& x, int a) const;
    virtual double getProb(int varIndex, double varVal, int a) const;
    virtual vector<GenDistribution*> createMarginals(const VarSet& x) const;
    virtual AbstractionSet* createJoint(list<int> queryVars, 
            const VarSet& evidence) const;

    void allocParameters();
    virtual void initParameters(VarSet* obs, int numObs);
    virtual void extendParameters();
    virtual void assignObservation(const VarSet& obs, int a, double p);
    virtual void updateParameters();
    virtual void saveParameters();
    virtual void restoreParameters();
    virtual void resetParameters();

    virtual istream& readParameters(istream& in);
    virtual ostream& writeParameters(ostream& out) const;

protected:
    void checkRelevance(int a, int var);
};

class AbsData {
public:
    double** varValProbs;
    double** savedVarValProbs;
    double** varValCounts;
    double*  totalVarCounts;
    bool*    relevant;
    VarSchema schema;

    AbsData(VarSchema s) :schema(s) { alloc(); }
    AbsData(const AbsData& other) 
        :varValProbs(0), savedVarValProbs(0), varValCounts(0), 
         totalVarCounts(0), relevant(0)
     { (*this) = other; }
    ~AbsData() { clear(); }

    void alloc()
    {
        int numVars = schema.getNumVars();
        varValProbs = new double*[numVars];
        varValCounts = new double*[numVars];
        savedVarValProbs = new double*[numVars];
        totalVarCounts = new double[numVars];
        relevant = new bool[numVars];

        // Initialize the new data
        for (int var = 0; var < numVars; var++) {

            int numVals = schema.getNumVals(var);
            varValProbs[var] = new double[numVals];
            varValCounts[var] = new double[numVals];
            savedVarValProbs[var] = new double[numVals];

            if (schema.isContinuous(var)) {
                // FIXME: not properly defining all values here...
                varValCounts[var][0] = 0.0;
                varValCounts[var][1] = INITIAL_WEIGHT;
                varValCounts[var][2] = 0.0;
                totalVarCounts[var] = INITIAL_WEIGHT;
            } else {
                for (int val = 0; val < numVals; val++) {
                    varValProbs[var][val] = 1.0/numVals;
                    varValCounts[var][val] = 1.0/numVals;
                    savedVarValProbs[var][val] = 0.0;
                }
                totalVarCounts[var] = 1.0;
            }
            relevant[var] = true;
        }
    }

    void clear()
    {
        int numVars = schema.getNumVars();
        for (int var = 0; var < numVars; var++) {
            delete [] varValProbs[var];
            delete [] varValCounts[var];
            delete [] savedVarValProbs[var];
        }
        delete [] varValProbs;
        delete [] varValCounts;
        delete [] savedVarValProbs;
        delete [] totalVarCounts;
        delete [] relevant;
    }

    const AbsData& operator=(const AbsData& other)
    {
        clear();
        schema = other.schema;
        alloc();

        for (int var = 0; var < schema.getNumVars(); var++) {
            for (int val = 0; val < schema.getNumVals(var); val++) {
                varValProbs[var][val] = other.varValProbs[var][val];
                varValCounts[var][val] = other.varValCounts[var][val];
                savedVarValProbs[var][val] = other.savedVarValProbs[var][val];
            }
            totalVarCounts[var] = other.totalVarCounts[var];
            relevant[var] = other.relevant[var];
        }
        return (*this);
    }
};
}
#endif // ndef VARABSTRACTIONSET_H
